#ifndef EIGEN_CONCAT_L_HPP
#define EIGEN_CONCAT_L_HPP

#include <Eigen/Core>
#include <type_traits>

namespace EigenCustom
{
using namespace Eigen;

constexpr Index dynamicOrSum( const Index& a, const Index& b ){
	return a == Dynamic || b == Dynamic ? Dynamic : a + b;
}

enum class Direction { horizontal, vertical };

template<Direction direction, typename Arg1, typename Arg2>
struct ConcatHelper {
	static_assert( std::is_same_v<
		typename Arg1::Scalar, typename Arg2::Scalar
	> );
	using Scalar = typename Arg1::Scalar;
	using D = Direction;
	static constexpr Index
		RowsAtCompileTime { direction == D::horizontal ?
			Arg1::RowsAtCompileTime :
			dynamicOrSum( Arg1::RowsAtCompileTime, Arg2::RowsAtCompileTime )
		},
		ColsAtCompileTime { direction == D::horizontal ?
			dynamicOrSum( Arg1::ColsAtCompileTime, Arg2::ColsAtCompileTime ) :
			Arg1::ColsAtCompileTime
		},
		MaxRowsAtCompileTime { direction == D::horizontal ?
			Arg1::MaxRowsAtCompileTime :
			dynamicOrSum( Arg1::MaxRowsAtCompileTime, Arg2::MaxRowsAtCompileTime )
		},
		MaxColsAtCompileTime { direction == D::horizontal ?
			dynamicOrSum( Arg1::MaxColsAtCompileTime, Arg2::MaxColsAtCompileTime ) :
			Arg1::MaxColsAtCompileTime
		};
	
	static_assert(
		(std::is_base_of_v<MatrixBase<Arg1>, Arg1> &&
		 std::is_base_of_v<MatrixBase<Arg2>, Arg2> ) ||
		(std::is_base_of_v<ArrayBase<Arg1>, Arg1> &&
		 std::is_base_of_v<ArrayBase<Arg2>, Arg2> )
	);
	using DenseType = std::conditional_t<
		std::is_base_of_v<MatrixBase<Arg1>, Arg1>,
		Matrix<
			Scalar,   RowsAtCompileTime,    ColsAtCompileTime,
			ColMajor, MaxRowsAtCompileTime, MaxColsAtCompileTime
		>,
		Array<
			Scalar,   RowsAtCompileTime,    ColsAtCompileTime,
			ColMajor, MaxRowsAtCompileTime, MaxColsAtCompileTime
		>
	>;
};
 
template<Direction direction, typename Arg1, typename Arg2>
class ConcatFunctor
{
	using Scalar = typename ConcatHelper<direction, Arg1, Arg2>::Scalar;
	const typename Arg1::Nested m_mat1;
	const typename Arg2::Nested m_mat2;

public:
	ConcatFunctor(const Arg1& arg1, const Arg2& arg2)
		: m_mat1(arg1), m_mat2(arg2)
	{}

	const Scalar operator() (Index row, Index col) const {
		if constexpr (direction == Direction::horizontal){
			if (col < m_mat1.cols())
				return m_mat1(row,col);
			return m_mat2(row, col - m_mat1.cols());
		} else {
			if (row < m_mat1.rows())
				return m_mat1(row,col);
			return m_mat2(row - m_mat1.rows(), col);
		}
	}
};

template<Direction direction, typename Arg1, typename Arg2>
using ConcatReturnType = CwiseNullaryOp<
	ConcatFunctor<direction,Arg1,Arg2>,
	typename ConcatHelper<direction,Arg1,Arg2>::DenseType
>;

template<Direction direction, typename Arg1, typename Arg2>
ConcatReturnType<direction, Arg1, Arg2>
concat(
	const Eigen::DenseBase<Arg1>& arg1,
	const Eigen::DenseBase<Arg2>& arg2
){
	using DenseType = typename ConcatHelper<direction,Arg1,Arg2>::DenseType;
	using D = Direction;
	return DenseType::NullaryExpr(
		direction == D::horizontal ? arg1.rows() : arg1.rows() + arg2.rows(),
		direction == D::horizontal ? arg1.cols() + arg2.cols() : arg1.cols(),
		ConcatFunctor<direction,Arg1,Arg2>( arg1.derived(), arg2.derived() )
	);
}

template<Direction direction, typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat(
	const Eigen::DenseBase<Arg1>& arg1,
	const Eigen::DenseBase<Arg2>& arg2,
	Ts&& ... rest
){
	return concat<direction>(
		concat<direction>(arg1, arg2),
		std::forward<Ts>(rest) ...
	);
}

template<typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat_horizontal(
	const Eigen::DenseBase<Arg1>& arg1,
	const Eigen::DenseBase<Arg2>& arg2,
	Ts&& ... rest
){
	return concat<Direction::horizontal>(
		arg1, arg2, std::forward<Ts>(rest) ...
	);
}

template<typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat_vertical(
	const Eigen::DenseBase<Arg1>& arg1,
	const Eigen::DenseBase<Arg2>& arg2,
	Ts&& ... rest
){
	return concat<Direction::vertical>(
		arg1, arg2, std::forward<Ts>(rest) ...
	);
}
	

} // namespace EigenCustom

namespace hms
{

using EigenCustom::concat;
using EigenCustom::concat_horizontal;
using EigenCustom::concat_vertical;

} // namespace hms



#endif
